import argparse
from random import random
from time import sleep
import os
import re
import copy
import numpy as np
import pandas as pd
from bins.PlannerLogger import PlannerLogger
from bins import vehicleParams
from bins.Map import Map
from src.Car_Iteration import CarPlanner
from src.motionplanner.Constrained_ILQR.scripts.arguments import add_arguments
from src.referenceLine import LaneChangeReferenceGenerator
from tools.drawAnimation import Animation


Logger = PlannerLogger()


class MotionPlanner:
    def __init__(self, args):
        self.args = args
        self.multi_car = {}
        self.map = Map()
        self.vehicle = vehicleParams.VehicleParams()
        self.count = 0


    def load_car(self, name, ref_line, v, x, y, v_, h):
        self.multi_car[name] = CarPlanner(copy.deepcopy(self.args), name)
        self.multi_car[name].update_refline(ref_line, v)
        self.multi_car[name].set_begin(x=x, y=y, v=v_, h=h)


    def set_car_restrain(self, car1_name, car2_name, distance):
        self.multi_car[car1_name].load_restrained_vehicle(car2_name, distance=distance)

    def update_car_state(self, name, speed, ref_line):
        generator = LaneChangeReferenceGenerator()
        initial_state = list(self.multi_car[name].current_state)
        initial_state[2] = speed
        ref_traj = generator.generate_reference(initial_state, ref_line)
        print(ref_traj)
        self.multi_car[name].update_reference(ref_traj, speed, ref_line)




    def count_all_trajectory(self, name: str=None):
        all_trajectory = {}
        for i in self.multi_car.keys():
            if i == name:
                continue
            if self.multi_car[i].state is None:
                continue
            all_trajectory[i] = self.multi_car[i].state.T[:, 1:self.args.horizon + 1]
        return all_trajectory

    def run_loop(self):
        self.count += 1
        for i in self.multi_car.keys():
            self.multi_car[i].update_npc_trajectory(self.count_all_trajectory(i))
            self.multi_car[i].update_control()
            rand = random()
            if rand < 0.3:
                rand = 1
            elif rand < 0.6:
                rand = 2
            elif rand < 0.9:
                rand = 3
            else:
                rand = 0
            self.multi_car[i].update_current_state(rand)
        Logger.record_(self.count)

    def load_all_trajectory(self, folder_path):
        files = os.listdir(folder_path)
        all_files = []
        for file_name in files:
            file_path = os.path.join(folder_path, file_name)
            if os.path.isfile(file_path):
                all_files.append(file_path)

    def save_all_trajectory(self, folder_path):
        if not os.path.exists(folder_path):
            os.mkdir(folder_path)
        folder_path = os.path.abspath(folder_path)
        for name in self.multi_car:
            self.save_trajectory(name, folder_path)

    def load_trajectory(self, file_name):
        """
        载入轨迹命名方式
        ./test1_2__3_4__70.csv
        test1为名称
        2为第几次迭代
        3_4为参考线y坐标值
        70为目标车速 km/h
        """
        file_name_ = os.path.basename(file_name)
        match = re.search(r'([^_]+)_(\d+)__(.+?)__([^.]+)\.csv', file_name_)
        if match:
            name = match.group(1)
            reference_line = match.group(3)
            reference_line_y = float(reference_line.replace('_', '.'))
            speed = float(match.group(4))
            self.load_car(name, reference_line_y, speed / 3.6)
            self.multi_car[name].load_trajectory(file_name)


    def save_trajectory(self, name, path):
        y = str(round(self.multi_car[name].desired_y, 1))
        y = y.replace('.', '_')
        v = str(int(self.multi_car[name].v * 3.6))
        file_name = f"{self.multi_car[name].name}_{self.count}__{y}__{v}.csv"
        if path[-1] != "/":
            path += "/"
        file_path = path + file_name
        np.savetxt(file_path, self.multi_car[name].state, delimiter=",", fmt="%.2f")

    def save_run_trajectory(self, folder):
        if not os.path.exists(folder):
            os.mkdir(folder)
        folder_path = os.path.abspath(folder)
        if folder_path[-1] != '/':
            folder_path += '/'
        for i in self.multi_car.keys():
            file_name = i + f"_{self.count}.csv"
            file_name = folder_path + file_name
            np.savetxt(file_name, np.array(self.multi_car[i].historical_track), delimiter=",", fmt="%.2f")







if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description='CARLA CILQR')
    add_arguments(argparser)
    args = argparser.parse_args()
    t = MotionPlanner(args)
    # t.load_all_trajectory("../trajectory/0_100/")
    for i in range(0, 10):
        name = "Car" + str(i)
        name_ = "Car" + str(i - 1)
        t.load_car(name=name, ref_line=3.4 / 2, v= 70 / 3.6, x=46 - i * 4, y=1.7, v_=0, h=0)
        if i != 0:
            t.set_car_restrain(car1_name=name, car2_name=name_, distance=10)

    for i in range(500):
        print(i)
        t.run_loop()

        # t.save_all_trajectory("../trajectory/test")
    t.save_run_trajectory("../trajectory/0_100/")
    t.save_all_trajectory("../trajectory/100alldata/")
    ani = Animation()
    for i in t.multi_car.keys():

        ani.load_trajectory(t.multi_car[i].historical_track)
        print(t.multi_car[i].historical_track)

    ani.draw()
